import os
import sys
from pathlib import Path

# 导入统一的环境配置
from env_setup import setup_environment, ENV_CONFIG

# 使用统一的环境配置
ENV_CONFIG = setup_environment()

try:
    from public.download_utils import download_file_with_hf
except ImportError:
    print(f"❌ 无法导入 download_utils")
    sys.exit(1)

CONFIG_FILE = ENV_CONFIG['CONFIG_DIR'] / "model_urls.conf"   # 配置文件名
SAVE_PATH = "downmodels"         # 下载保存目录
DEFAULT_FILENAME = ""            # 可选：默认使用 URL 提取的文件名

def read_urls_from_file(config_file):
    urls = []
    if not os.path.exists(config_file):
        print(f"❌ 找不到配置文件: {config_file}")
        return urls

    with open(config_file, "r", encoding="utf-8-sig") as f:  # utf-8-sig 可处理 BOM
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):  # 跳过空行和注释
                continue
            # 可选：移除行内注释（# 之后的内容）
            if "#" in line:
                line = line.split("#", 1)[0].strip()
            if line:
                urls.append(line)
    return urls

from concurrent.futures import ThreadPoolExecutor, as_completed

def download_task(url, idx, total):
    print(f"\n=== [{idx}/{total}] 开始下载: {url}")
    success, result = download_file_with_hf(url=url, filename=DEFAULT_FILENAME, save_path=SAVE_PATH)
    if success:
        print(f"✅ 文件已保存: {result}")
        return url, True
    else:
        print(f"❌ 下载失败: {url}, 错误: {result}")
        return url, False

def main():
    urls = read_urls_from_file(CONFIG_FILE)
    if not urls:
        print("⚠️ 配置文件中没有有效 URL")
        return

    # 使用4个线程并发下载
    with ThreadPoolExecutor(max_workers=4) as executor:
        futures = [
            executor.submit(download_task, url, idx+1, len(urls))
            for idx, url in enumerate(urls)
        ]
        
        for future in as_completed(futures):
            url, success = future.result()

if __name__ == "__main__":
    main()
